<!DOCTYPE html>
<html>
<head>
    <meta charset="utf-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0, user-scalable=yes">
    <title>WebGPU Simple Compute Shader</title>
    <style>
        @import url(res/webgpu-lesson.css);
    </style>
</head>
<body>
    <canvas></canvas>
</body>
<script type="module">
    async function main() {
        const adapter = await navigator.gpu?.requestAdapter();
        const device = await adapter?.requestDevice();
        if (!device) {
            fail('need a browser that supports WebGPU');
            return;
        }

        // data 变量，类型是storage，它既能被读取，也能被写入
        // 在 bindGroup 0（@group(0)）中的绑定位置 0（binding(0)）上设置这个data数组
        const module = device.createShaderModule({
            label: 'doubling compute module',
            code: `
                @group(0) @binding(0) var<storage, read_write> data: array<f32>;

                @compute @workgroup_size(1) fn computeSomething(
                    @builtin(global_invocation_id) id: vec3u
                ) {
                    let i = id.x;
                    data[i] = data[i] * 2.0;
                }
            `,
        });

        const pipeline = device.createComputePipeline({
            label: 'doubling compute pipeline',
            layout: 'auto',
            compute: {
                module,
            },
        });
        
        const input = new Float32Array([1, 3, 5]);
        // 在 GPU 上创建一个缓冲区，并将CPU数据复制到缓冲区中
        const workBuffer = device.createBuffer({
            label: 'work buffer',
            size: input.byteLength, // 缓冲区的大小, 以字节为单位
            // GPUBufferUsage.STORAGE 来表示希望将此缓冲区用作storage用途。这样就可以与着色器中的 var<storage,...> 兼容
            // GPUBufferUsage.COPY_SRC 表示希望将此缓冲区用作数据源，希望能从该缓冲区复制数据
            // GPUBufferUsage.COPY_DST 表示希望将此缓冲区用作数据目的地，希望能将数据复制到该缓冲区
            usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST,
        });

        device.queue.writeBuffer(workBuffer, 0, input);

        // 不能直接从 JavaScript 中读取 WebGPU 缓冲区的内容
        // 需要另一个缓冲区。运行计算后，把上面的缓冲区复制到这个结果缓冲区，并设置其标志，以便进行映射
        const resultBuffer = device.createBuffer({
            label: 'result buffer',
            size: input.byteLength,
            //MAP_READ 表示能够映射该缓冲区以读取数据
            usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST,
        });

        // 为了告诉着色器希望它在哪个缓冲区上工作，需要 bindGroup
        const bindGroup = device.createBindGroup({
            label: 'bindGroup for work buffer',
            // 0 对应着着色器中的 @group(0)
            layout: pipeline.getBindGroupLayout(0),
            entries: [
            // {binding：0 ... entry对应着着色器中的 @group(0) @binding(0)
                { binding: 0, resource: { buffer: workBuffer } },
            ],
        });

        const encoder = device.createCommandEncoder({
            label: 'doubling encoder',
        });

        const pass = encoder.beginComputePass({
            label: 'doubling compute pass',
        });

        pass.setPipeline(pipeline); 
        // 0 对应着着色器中的 @group(0)
        pass.setBindGroup(0, bindGroup);
        pass.dispatchWorkgroups(input.length);
        pass.end();

        encoder.copyBufferToBuffer(workBuffer, 0, resultBuffer, 0, resultBuffer.size);

        const commandBuffer = encoder.finish();
        device.queue.submit([commandBuffer]);   

        // 要映射结果缓冲区，需要调用 mapAsync 并需要使用await等待其完成
        await resultBuffer.mapAsync(GPUMapMode.READ);
        // resultBuffer.getMappedRange()（无参数），返回整个缓冲区的 ArrayBuffer
        // getMappedRange 返回的 ArrayBuffer 仅在调用 unmap 之前有效。在unmap之后，其长度将被设置为 0，数据也不再可访问
        const result = new Float32Array(resultBuffer.getMappedRange().slice());
        // 完成映射后，需要调用 unmap() 来释放缓冲区
        resultBuffer.unmap();

        console.log('input: ', input);
        console.log('result: ', result);
    }


    function fail(msg) {
        // eslint-disable-next-line no-alert
        alert(msg);
    }

    main();
</script>   
</html>